#######################################################################
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# Openfuyao is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#     http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
# PURPOSE.
# See the Mulan PSL v2 for more details.
#######################################################################

#!/bin/bash

# 获取当前脚本所在目录的绝对路径
SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)

# 默认参数
TEST_RUN_TIMES=1
TEST_TYPE="default"
NGINX_SERVER_YAML="${SCRIPT_DIR}/test-deployment-nginx-server.yaml"
NGINX_SERVICE_YAML="${SCRIPT_DIR}/test-service-nginx-server.yaml"
BENCHMARK_YAML="${SCRIPT_DIR}/test-job-wrk-client.yaml"
NGINX_CPU_REQUEST="1"
NGINX_CPU_LIMIT="1"
BENCHMARK_CPU_REQUEST="2"
BENCHMARK_CPU_LIMIT="2"
NAMESPACE="default"
WAIT_TIMEOUT=300  # 等待Nginx就绪的超时时间（秒）
BENCHMARK_SCRIPT="${SCRIPT_DIR}/auto-wrk-client.sh"  # 性能测试脚本路径
CLEANUP=true      # 是否清理Nginx资源
CONCURRENCY="200" # 默认并发数
DURATION="300s"   # 默认持续时间
CPU_TEST_BINARY="${SCRIPT_DIR}/../cputest.arm"  # 上一级目录的CPU压力测试工具
STRESS_TEST=true  # 是否执行加压测试
STRESS_COMMAND="" # 压力测试命令（通过参数传入）

# 帮助信息
function show_help() {
    echo "用法: $0 -c <并发数> -d <持续时间> [选项]"
    echo "选项:"
    echo "  -c <并发数>          并发连接数 (默认: 200)"
    echo "  -d <持续时间>        压测持续时间 (如: 300s, 5m, 1m30s) (默认: 300s)"
    echo "  -n <次数>            运行测试的次数 (默认: 1)"
    echo "  -t <类型>            测试类型，用于日志文件名 (默认: default)"
    echo "  -yServer <路径>      Nginx Server部署YAML文件路径 (默认: test-deployment-nginx-server.yaml)"
    echo "  -yService <路径>     Nginx Service YAML文件路径 (默认: test-service-nginx-server.yaml)"
    echo "  -yClient <路径>      Wrk Client部署YAML文件路径 (默认: test-job-wrk-client.yaml)"
    echo "  -reqServer <CPU>     Nginx Server的CPU requests (默认: 1)"
    echo "  -limitServer <CPU>   Nginx Server的CPU limits (默认: 1)"
    echo "  -reqClient <CPU>     Wrk Client的CPU requests (默认: 2)"
    echo "  -limitClient <CPU>   Wrk Client的CPU limits (默认: 2)"
    echo "  --no-cleanup        不清理Nginx资源"
    echo "  -stressCommand <命令> 加压测试命令（例如: '../cputest.arm 0 0:50,3-127:50')"
    echo "  --no-stress         不执行加压测试"
    echo "  -h                  显示此帮助信息"
    exit 1
}

# 等待Nginx服务就绪
function wait_for_nginx() {
    echo "等待Nginx服务就绪..."
    local start_time=$(date +%s)
    
    while true; do
        # 检查就绪Pod数量
        READY_PODS=$(kubectl get pods -n "$NAMESPACE" -l app=nginx -o=jsonpath='{range .items[*]}{@.status.conditions[?(@.type=="Ready")].status}{"\n"}{end}' | grep -c "True")
        
        if [[ "$READY_PODS" -gt 0 ]]; then
            echo "Nginx服务已就绪!"
            return 0
        fi
        
        # 检查是否超时
        current_time=$(date +%s)
        elapsed_time=$((current_time - start_time))
        
        if [[ $elapsed_time -gt $WAIT_TIMEOUT ]]; then
            echo "错误: 等待Nginx服务超时 ($WAIT_TIMEOUT 秒)"
            return 1
        fi
        
        echo "等待Nginx服务就绪中... ($elapsed_time/$WAIT_TIMEOUT秒)"
        sleep 5
    done
}

# 清理Nginx资源
function cleanup_nginx() {
    echo "开始清理Nginx资源..."
    
    # 尝试按标签删除资源
    kubectl delete deployments -l app=nginx-server-deployment -n "$NAMESPACE" 2>/dev/null
    kubectl delete services -l app=nginx-server-service -n "$NAMESPACE" 2>/dev/null
    
    # 等待资源删除完成
    local start_time=$(date +%s)
    while true; do
        POD_COUNT=$(kubectl get pods -n "$NAMESPACE" -l app=nginx -o jsonpath='{.items.length}' 2>/dev/null || echo 0)
        
        if [[ "$POD_COUNT" -eq 0 ]]; then
            echo "Nginx资源清理完成!"
            return 0
        fi
        
        current_time=$(date +%s)
        elapsed_time=$((current_time - start_time))
        
        if [[ $elapsed_time -gt 120 ]]; then
            echo "警告: Nginx资源清理超时，但继续执行..."
            return 1
        fi
        
        echo "等待Nginx资源清理中... ($elapsed_time/120秒)"
        sleep 5
    done
}

# 执行基准测试
function run_benchmark() {
    local test_type="$1"
    local nginx_url="$2"
    local test_label="${test_type}"
    
    echo "开始${test_label}性能测试..."
    
    # 构建性能测试命令
    BENCHMARK_CMD="$BENCHMARK_SCRIPT -n $TEST_RUN_TIMES -t $test_type -c $CONCURRENCY -d $DURATION -y $BENCHMARK_YAML -u $nginx_url"
    
    if [[ -n "$BENCHMARK_CPU_REQUEST" ]]; then
        BENCHMARK_CMD="$BENCHMARK_CMD -req $BENCHMARK_CPU_REQUEST"
    fi
    
    if [[ -n "$BENCHMARK_CPU_LIMIT" ]]; then
        BENCHMARK_CMD="$BENCHMARK_CMD -limit $BENCHMARK_CPU_LIMIT"
    fi
    
    echo "${test_label}测试命令: $BENCHMARK_CMD"
    
    # 执行性能测试并捕获退出码
    start_time=$(date +%s)
    eval "$BENCHMARK_CMD"
    TEST_EXIT_CODE=$?
    end_time=$(date +%s)
    
    if [[ $TEST_EXIT_CODE -eq 0 ]]; then
        echo "${test_label}性能测试成功完成 (耗时: $((end_time - start_time))秒)"
        return 0
    else
        echo "${test_label}性能测试失败，退出码: $TEST_EXIT_CODE"
        return 1
    fi
}

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        -c)
            CONCURRENCY="$2"
            shift 2
            ;;
        -d)
            DURATION="$2"
            shift 2
            ;;
        -n)
            TEST_RUN_TIMES="$2"
            shift 2
            ;;
        -t)
            TEST_TYPE="$2"
            shift 2
            ;;
        -yServer)
            NGINX_SERVER_YAML="$2"
            shift 2
            ;;
        -yService)
            NGINX_SERVICE_YAML="$2"
            shift 2
            ;;
        -yClient)
            BENCHMARK_YAML="$2"
            shift 2
            ;;
        -reqServer)
            NGINX_CPU_REQUEST="$2"
            shift 2
            ;;
        -limitServer)
            NGINX_CPU_LIMIT="$2"
            shift 2
            ;;
        -reqClient)
            BENCHMARK_CPU_REQUEST="$2"
            shift 2
            ;;
        -limitClient)
            BENCHMARK_CPU_LIMIT="$2"
            shift 2
            ;;
        --no-cleanup)
            CLEANUP=false
            shift
            ;;
        -stressCommand)
            STRESS_COMMAND="$2"
            shift 2
            ;;
        --no-stress)
            STRESS_TEST=false
            shift
            ;;
        -h)
            show_help
            ;;
        *)
            echo "未知参数: $1" >&2
            show_help
            ;;
    esac
done

# 验证参数
if ! [[ "$CONCURRENCY" =~ ^[0-9]+$ ]]; then
    echo "错误: 并发数必须是数字" >&2
    show_help
fi

if ! [[ "$DURATION" =~ ^([0-9]+[sm]?)+$ ]]; then
    echo "错误: 持续时间格式不正确 (示例: 120s, 5m, 1m30s)" >&2
    show_help
fi

if ! [[ "$TEST_RUN_TIMES" =~ ^[0-9]+$ ]]; then
    echo "错误: 运行次数必须是数字" >&2
    show_help
fi

# 验证CPU参数格式
for param in "$NGINX_CPU_REQUEST" "$NGINX_CPU_LIMIT" "$BENCHMARK_CPU_REQUEST" "$BENCHMARK_CPU_LIMIT"; do
    if ! [[ "$param" =~ ^([0-9]+m|[0-9]+(\.[0-9]+)?)$ ]]; then
        echo "错误: CPU参数格式不正确 (示例: 2, 1.5, 500m)" >&2
        show_help
    fi
done

# 检查YAML文件是否存在
if [[ ! -f "$NGINX_SERVER_YAML" ]]; then
    echo "错误: Nginx Server YAML文件不存在: $NGINX_SERVER_YAML"
    exit 1
fi

if [[ ! -f "$NGINX_SERVICE_YAML" ]]; then
    echo "错误: Nginx Service YAML文件不存在: $NGINX_SERVICE_YAML"
    exit 1
fi

if [[ ! -f "$BENCHMARK_YAML" ]]; then
    echo "错误: Wrk Client YAML文件不存在: $BENCHMARK_YAML"
    exit 1
fi

# 检查CPU压力测试工具是否存在
if [[ "$STRESS_TEST" == "true" ]]; then
    if [[ ! -f "$CPU_TEST_BINARY" ]]; then
        echo "错误: CPU压力测试工具不存在: $CPU_TEST_BINARY"
        echo "请确保 $CPU_TEST_BINARY 存在且可执行，或使用 --no-stress 选项跳过加压测试"
        exit 1
    elif [[ ! -x "$CPU_TEST_BINARY" ]]; then
        # 修改CPU压力测试工具权限
        chmod +x "$CPU_TEST_BINARY"
    fi
fi

# 检查性能测试脚本是否存在
if [[ ! -f "$BENCHMARK_SCRIPT" ]]; then
    echo "错误: 性能测试脚本不存在: $BENCHMARK_SCRIPT"
    exit 1
elif [[ ! -x "$BENCHMARK_SCRIPT" ]]; then
    # 修改性能测试脚本权限
    chmod +x "$BENCHMARK_SCRIPT"
fi

echo "开始部署Nginx服务..."

# 创建临时YAML文件（当前目录）
TEMP_NGINX_YAML="${SCRIPT_DIR}/nginx-server-$(date +%s).yaml"
TEMP_SERVICE_YAML="${SCRIPT_DIR}/nginx-service-$(date +%s).yaml"

# 复制并修改Nginx Server的YAML
cp "$NGINX_SERVER_YAML" "$TEMP_NGINX_YAML"
# 使用yq修改requests.cpu (兼容语法)
yq -i -y ".spec.template.spec.containers[0].resources.requests.cpu = \"${NGINX_CPU_REQUEST}\"" "$TEMP_NGINX_YAML"
# 使用yq修改limits.cpu (兼容语法)
yq -i -y ".spec.template.spec.containers[0].resources.limits.cpu = \"${NGINX_CPU_LIMIT}\"" "$TEMP_NGINX_YAML"

# 复制Nginx Service的YAML（无需修改）
cp "$NGINX_SERVICE_YAML" "$TEMP_SERVICE_YAML"

# 应用Nginx部署和服务
kubectl apply -f "$TEMP_SERVICE_YAML" -n "$NAMESPACE"
kubectl apply -f "$TEMP_NGINX_YAML" -n "$NAMESPACE"

# 等待Nginx服务就绪
if ! wait_for_nginx; then
    echo "部署Nginx服务失败!"
    # 清理临时文件和资源
    rm -f "$TEMP_NGINX_YAML" "$TEMP_SERVICE_YAML"
    if [[ "$CLEANUP" == "true" ]]; then
        cleanup_nginx
    fi
    exit 1
fi

# 获取Nginx服务信息
NGINX_SERVICE_IP=$(kubectl get svc nginx-service -n "$NAMESPACE" -o jsonpath='{.spec.clusterIP}')
NGINX_SERVICE_PORT=$(kubectl get svc nginx-service -n "$NAMESPACE" -o jsonpath='{.spec.ports[0].port}')
NGINX_URL="http://${NGINX_SERVICE_IP}:${NGINX_SERVICE_PORT}"
echo "Nginx服务地址: $NGINX_URL"

# 执行基准测试
echo "========================================="
echo "开始基准性能测试"
echo "========================================="
run_benchmark "$TEST_TYPE" "$NGINX_URL"
BASELINE_EXIT_CODE=$?

# 执行加压测试（如果启用且基准测试成功）
if [[ "$STRESS_TEST" == "true" && $BASELINE_EXIT_CODE -eq 0 && -n "$STRESS_COMMAND" ]]; then
    echo "========================================="
    echo "准备执行加压性能测试"
    echo "加压命令: $STRESS_COMMAND"
    echo "========================================="
    
    # 启动压力测试（后台运行）
    eval "$STRESS_COMMAND &"
    STRESS_PID=$!
    echo "加压测试已启动，PID: $STRESS_PID"
    
    # 等待压力测试生效（可根据需求调整等待时间）
    sleep 240
    
    # 执行压力下的性能测试
    run_benchmark "${TEST_TYPE}-stress" "$NGINX_URL"
    STRESS_EXIT_CODE=$?
    
    # 停止压力测试
    echo "停止加压测试..."
    kill -TERM "$STRESS_PID" 2>/dev/null
    wait "$STRESS_PID" 2>/dev/null
    
    [[ $STRESS_EXIT_CODE -ne 0 ]] && BASELINE_EXIT_CODE=$STRESS_EXIT_CODE
elif [[ "$STRESS_TEST" == "true" ]]; then
    echo "错误: 未提供加压命令 (-stressCommand)，跳过加压测试"
fi

# 清理临时文件
rm -f "$TEMP_NGINX_YAML" "$TEMP_SERVICE_YAML"

# 清理Nginx资源
if [[ "$CLEANUP" == "true" ]]; then
    cleanup_nginx
    CLEANUP_EXIT_CODE=$?
    
    # 如果清理失败，更新最终退出码
    if [[ $CLEANUP_EXIT_CODE -ne 0 && $BASELINE_EXIT_CODE -eq 0 ]]; then
        BASELINE_EXIT_CODE=$CLEANUP_EXIT_CODE
    fi
fi

# 输出最终结果
echo "========================================="
if [[ $BASELINE_EXIT_CODE -eq 0 ]]; then
    echo "所有测试已成功完成!"
else
    echo "测试执行失败，退出码: $BASELINE_EXIT_CODE"
fi
echo "========================================="

exit $BASELINE_EXIT_CODE